#include "cast_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling {
    const uint32_t BLOCK_SIZE = 32;
    const uint32_t BUFFER_NUM = 2;
    const uint32_t VAL_ZRRO = 0;
    int64_t diag_Val = VAL_ZRRO;
    int32_t tilingKeyBase = 0;
    int32_t offset = 1; 
    
    
    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {
        TilingData tiling;
        uint64_t ubSize;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
        auto coreNum = ascendcPlatform.GetCoreNum();
        
        // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
        uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
        uint32_t typeLength = 0;
        ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
        uint32_t inputLength = inputNum * typeLength;
        uint32_t inputBytes = inputLength / inputNum;

        // There are a total of 3 shared UB spaces in the input and output. If it's int8, there are 2 more TBUFs
        uint32_t ubDataNumber = (inputBytes == 1) ? 5 : 5;
        // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
        uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
        uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;

        // Input data for 32B alignment
        uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
        // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
        coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
        coreNum = (coreNum >= 1) ? coreNum : 1;
        uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;
        uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;
        
        // Small chunks are calculated and sliced several times using the number of data on each core
        uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
        uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
        // Tail block calculation for small chunks of data
        uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
        smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
        
        // The total length of a large block of data is 32B larger than that of a small block of data
        everyCoreInputBlockNum += 1;
        uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
        uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
        uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
        bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;
        
        tiling.set_smallCoreDataNum(smallCoreDataNum);
        tiling.set_bigCoreDataNum(bigCoreDataNum);
        tiling.set_tileDataNum(tileDataNum);
        tiling.set_smallTailDataNum(smallTailDataNum);
        tiling.set_bigTailDataNum(bigTailDataNum);
        tiling.set_finalSmallTileNum(finalSmallTileNum);
        tiling.set_finalBigTileNum(finalBigTileNum);
        tiling.set_tailBlockNum(tailBlockNum);
        
        context->SetBlockDim(coreNum);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;


        //根据p值和数据类型共同决定tiling key,第一个标识符为输入数据类型，第二个标识符为输出数据类型
            // TilingKey 编码规则:
            //input
            // half        -> 00   
            // float       -> 10 
            // bfloat16_t  -> 20
            // int4b_t     -> 30
            // uint8_t     -> 40
            // int8_t      -> 50 
            // int16_t     -> 60
            // int32_t     -> 70 
            // int64_t     -> 80
            
            //input
            // half        -> 0   
            // float       -> 1
            // bfloat16_t  -> 2
            // int4b_t     -> 3
            // uint8_t     -> 4
            // int8_t      -> 5
            // int16_t     -> 6
            // int32_t     -> 7
            // int64_t     -> 8
        const auto inputDataType = context->GetInputDesc(0)->GetDataType();
        const auto outputDataType = context->GetInputDesc(1)->GetDataType();

        switch (inputDataType) {
            case ge::DT_FLOAT16: tilingKeyBase = 0;  break; 
            case ge::DT_FLOAT:   tilingKeyBase = 10; break; 
            case ge::DT_UINT8:   tilingKeyBase = 20; break;
            case ge::DT_INT8:    tilingKeyBase = 30; break; 
            case ge::DT_INT16:   tilingKeyBase = 40; break; 
            case ge::DT_INT32:   tilingKeyBase = 50; break; 
            case ge::DT_INT64:   tilingKeyBase = 60; break; 
            default: return ge::GRAPH_FAILED;
        }

        /* ---------- 输出类型 ---------- */
        switch (outputDataType) {
            case ge::DT_FLOAT16: offset = 0; break;
            case ge::DT_FLOAT:   offset = 1; break;
            case ge::DT_UINT8:   offset = 2; break; 
            case ge::DT_INT8:    offset = 3; break; 
            case ge::DT_INT16:   offset = 4; break; 
            case ge::DT_INT32:   offset = 5; break; 
            case ge::DT_INT64:   offset = 6; break; 
            default: return ge::GRAPH_FAILED;
        }

        
        context->SetTilingKey(tilingKeyBase + offset);
        return ge::GRAPH_SUCCESS;
    }
} // namespace optiling

namespace ge {
static graphStatus InferShape(gert::InferShapeContext *context)
{
    const gert::Shape *x1_shape = context->GetInputShape(0);
    gert::Shape *y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}

static graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    const auto inputDataType = context->GetInputDataType(0);
    const auto outputDataType = context->GetInputDataType(1);
    context->SetOutputDataType(0, outputDataType);
        
    return ge::GRAPH_SUCCESS;
}
} // namespace ge




namespace ops {
class CastCustom : public OpDef {
public:
    explicit CastCustom(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({
                ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, 
                ge::DT_FLOAT,   ge::DT_FLOAT,   ge::DT_FLOAT,   ge::DT_FLOAT,   ge::DT_FLOAT,   ge::DT_FLOAT, 
                ge::DT_UINT8,   ge::DT_UINT8,   ge::DT_UINT8,   ge::DT_UINT8,   ge::DT_UINT8,   ge::DT_UINT8, 
                ge::DT_INT8,    ge::DT_INT8,    ge::DT_INT8,    ge::DT_INT8,    ge::DT_INT8,    ge::DT_INT8,    
                ge::DT_INT16,   ge::DT_INT16,   ge::DT_INT16,   ge::DT_INT16,   ge::DT_INT16,   ge::DT_INT16,   
                ge::DT_INT32,   ge::DT_INT32,   ge::DT_INT32,   ge::DT_INT32,   ge::DT_INT32,   ge::DT_INT32,   
                ge::DT_INT64,   ge::DT_INT64,   ge::DT_INT64,   ge::DT_INT64,   ge::DT_INT64,   ge::DT_INT64
            })
            .Format({
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND
            })
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({
                ge::DT_FLOAT,   ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32
            })
            .Format({
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND
            })
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({
                ge::DT_FLOAT,   ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT32, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT64,
                ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_UINT8, ge::DT_INT8,  ge::DT_INT16, ge::DT_INT32
            })
            .Format({
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, 
                ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND
            })
            .UnknownShapeFormat({ge::FORMAT_ND});
        
        
        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend910b");
    }
};
OP_ADD(CastCustom);
} // namespace ops
